package main;

import static crypt.EncryptUtil.*;

import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.ObjectOutputStream;
import java.security.InvalidKeyException;
import java.security.Key;
import java.security.KeyStore;
import java.security.KeyStoreException;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.SecureRandom;
import java.security.Signature;
import java.security.SignatureException;
import java.security.UnrecoverableKeyException;
import java.security.cert.CertificateException;
import java.security.spec.InvalidKeySpecException;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.CipherOutputStream;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.KeyGenerator;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.SecretKey;

import crypt.CryptoVariables;
import crypt.Configuration;
import crypt.EncryptUtil;

/**
 * WriteEncryption main
 * Input: File_To_Encrypt, Configuration_File_Path, Keystore_Password 
 * Output: encrypt <File_To_Encrypt> and write it and other variables to <Configuration_File_Path>
 */
public class WriteEncryption {

	/**
	 * WriteEncryption main
	 */
	public static void main(String[] args) throws NoSuchAlgorithmException,
			NoSuchPaddingException, InvalidKeyException, KeyStoreException,
			CertificateException, IOException, UnrecoverableKeyException,
			InvalidKeySpecException, IllegalBlockSizeException,
			BadPaddingException, SignatureException {

		if (args.length != 3) {
			System.out.println("Write " + "<plain_data_file> "
					+ "<configuration_file_path> " + "<Keystore_password>");
			return;
		}

		String plainDataFilePath = args[0];
		String configurationFilePath = args[1];
		String keystorePassword = args[2];
		Configuration config = new Configuration(keystorePassword);

		// Initialize the secured random
		SecureRandom random = createSecureRandom();

		// Initialize and store the key
		SecretKey cipherKey = getRandomCipherKey(random);


		// Sign and output the data
		byte[] data = EncryptUtil.readFully(plainDataFilePath);

        // Calculate the signature of the data
        byte[] signature = signData(data, keystorePassword);


		byte[] iv = writeDataAndSig(config.getEncryptedFilePath(), cipherKey,
				data, keystorePassword);

		// Output the cryptographic variables
		writeCryptoVariables(config.getCryptVarFilePath(), cipherKey, iv, keystorePassword);

		// Output the configuration
		writeConfiguration(configurationFilePath, config, signature);
	}

	/**
	 * Write the configuration for the deciphering process to the FS
	 * 
	 * @param configurationFilePath
	 *            The path to write the configuration to
	 * @param config
	 *            The DecryptionConfiguration to write
	 */
	private static void writeConfiguration(String configurationFilePath,
			Configuration config, byte[] signature) throws FileNotFoundException,
			IOException {

		System.out.println("Writing the configruation file to "
				+ configurationFilePath);
		ObjectOutputStream oos = null;
		try {
			FileOutputStream fos = new FileOutputStream(configurationFilePath);
			oos = new ObjectOutputStream(fos);
			oos.writeObject(config);
            oos.writeObject(signature);
			oos.flush();
		} finally {
			oos.close();
		}
	}

	/**
	 * Write the plain data and it's signature using CipherOutputStream
	 * 
	 * @param outputFilePath
	 *            The path of the PS to write the encrypted data
	 * @param cipherKey
	 *            The SecretKey to use for the encryption
	 * @param plainData
	 *            The plain data to write
	 * @param keystorePassword
	 *            The password to the keystore
	 * @return the IV used for the encryption
	 */
	private static byte[] writeDataAndSig(String outputFilePath,
			SecretKey cipherKey, byte[] plainData, String keystorePassword) throws IOException,
			InvalidKeyException, NoSuchAlgorithmException,
			NoSuchPaddingException, UnrecoverableKeyException,
			KeyStoreException, CertificateException, SignatureException {

		// Generate digital signature
		byte[] signatureData = signData(plainData, keystorePassword);

		// Output the encrypted bytes using CipherOutputStream
		Cipher cipher = createDataCipher(cipherKey);
		ObjectOutputStream cipherOos = null;
		try {
			FileOutputStream fos = new FileOutputStream(outputFilePath);
			CipherOutputStream cos = new CipherOutputStream(fos, cipher);
			cipherOos = new ObjectOutputStream(cos);

			System.out.println("Writing data: " + new String(plainData));
			cipherOos.writeObject(plainData);

			System.out.println("Writing signature: "
					+ EncryptUtil.getHexDump(signatureData)); //
			cipherOos.writeObject(signatureData);

			return cipher.getIV();
		} finally {
			cipherOos.close();
		}

	}

	/**
	 * Write the cryptogaphic variables to argument cryptoVarFilePath
	 * 
	 * @param cryptoVarFilePath
	 *            The path on the FS to write the cryptographic variables to
	 * @param keystorePassword
	 *            The password to the keystore
	 * @param iv
	 *            The IV used for the encryption of the data
	 */
	private static void writeCryptoVariables(String cryptoVarFilePath,
			SecretKey secretKey, byte[] iv, String keystorePassword) throws IOException,
			InvalidKeyException, UnrecoverableKeyException,
			NoSuchAlgorithmException, KeyStoreException, CertificateException,
			NoSuchPaddingException, IllegalBlockSizeException,
			BadPaddingException {

		System.out.println("Writing secret key: "
				+ EncryptUtil.getHexDump(secretKey.getEncoded()));
		System.out.println("Writing IV: " + EncryptUtil.getHexDump(iv));

		// Create the CryptoVariables instance

		KeyStore keyStore = EncryptUtil.getKeyStore(keystorePassword);
		Key publicKey = keyStore.getCertificate(RSA_ALIAS).getPublicKey();
		CryptoVariables cryptoVariables = new CryptoVariables(secretKey, iv,
				publicKey);

		ObjectOutputStream oos = null;
		try {
			FileOutputStream fos = new FileOutputStream(cryptoVarFilePath);
			oos = new ObjectOutputStream(fos);
			oos.writeObject(cryptoVariables);
			oos.flush();
		} finally {
			oos.close();
		}
	}

	/**
	 * Create the Cipher to use for the encryption
	 * 
	 * @param cipherKey
	 *            The SecretKey to encrypt with
	 * @return the created Cipher instance
	 */
	private static Cipher createDataCipher(SecretKey cipherKey)
			throws NoSuchAlgorithmException, NoSuchPaddingException,
			InvalidKeyException, FileNotFoundException {

		Cipher cipher = Cipher.getInstance(DATA_CIPHER_ALGORITHM);
		cipher.init(Cipher.ENCRYPT_MODE, cipherKey);
		return cipher;
	}

	/**
	 * Signs the argument data
	 * 
	 * @param plainData
	 *            The data to sign
	 * @return the signature over the data
	 */
	private static byte[] signData(byte[] plainData, String keystorePassword)
			throws FileNotFoundException, KeyStoreException, IOException,
			NoSuchAlgorithmException, CertificateException,
			UnrecoverableKeyException, InvalidKeyException, SignatureException {

		KeyStore keyStore = EncryptUtil.getKeyStore(keystorePassword);
		PrivateKey privateKey = (PrivateKey) keyStore.getKey(RSA_ALIAS,
				keystorePassword.toCharArray());
		Signature signature = Signature.getInstance(SIGNATURE_ALGORITHM);
		signature.initSign(privateKey);
		signature.update(plainData);
		return signature.sign();
	}

	/**
	 * Create a random secret key
	 * 
	 * @param random
	 *            The SecureRandom to use
	 * @return the created SecretKey
	 */
	private static SecretKey getRandomCipherKey(SecureRandom random)
			throws NoSuchAlgorithmException {

		KeyGenerator keyGenerator = KeyGenerator
				.getInstance(KEY_GENERATOR_ALGORITHM);
		keyGenerator.init(ENCRYPTION_KEY_SIZE, random);
		SecretKey key = keyGenerator.generateKey();
		return key;
	}

	/**
	 * Create and initialize the secure random.
	 * 
	 * @return a SecureRandom instance which is initialised using the current
	 *         time
	 */
	private static SecureRandom createSecureRandom()
			throws NoSuchAlgorithmException
    {
		SecureRandom random = SecureRandom.getInstance(SECURE_RANDOM_ALGORITHM);
		random.setSeed(System.currentTimeMillis());
		return random;
	}

}
